from .base_strategy import BaseStrategy
import backtrader as bt
import backtrader.indicators as btind
import indicator.MyTTIndicator as Mytt
import indicator.myzhibiao as mmm


###
# 马丁交易策略
# 当价格偏离均线过多，且，出现了上上影线时，则意味着当前价格处于超卖。可以进场卖出
# 
# #
class MyStrategy(BaseStrategy):
    params = (
        ('initial_risk', 0.30),  # 初始仓位10%
        ('add_risk', 0.1),      # 每次加仓4%
        ('max_adds', 5),         # 最大加仓次数
        ('trigger', 0.004),      # 每次下跌0.4%触发加仓
        ('stop_loss', 0.05),    # 2.2%止损
        ('won_loss', 0.02),    # 2%止盈
    )

 
    def __init__(self):
        super().__init__()

        # 跟踪订单和仓位状态
        self.order = None
        self.buy_price = 0.0
        self.buy_size = 0.0
        self.add_count = 0
        self.stop_price = 0.0
        
        # 指标
        self.sma = btind.SimpleMovingAverage(self.data.close, period=15)

    def notify_order(self, order):
        super().notify_order(order)

    def notify_trade(self, trade):
        super().notify_trade(trade)

    def next(self):
        super().next()
        
        # 检查是否有未完成的订单
        # if self.order:
        #     return
            
        # 检查是否已经有持仓
        if not self.position:
            # 初始买入条件
            if self.data.close[0] > self.sma[0]:  # 简单均线过滤
                self.initial_buy()
        else:
            # 检查止损
            current_loss = (self.data.close[0] - self.buy_price) / self.buy_price
            if current_loss <= -self.p.stop_loss:
                self.close_all_positions()
                return
            
            # 检查止盈
            if current_loss>=self.p.won_loss and self.data.close[0]<self.sma[0]:
                self.close_all_positions('止盈')
                return

            # 检查加仓条件
            current_drop = (self.data.close[0] - self.buy_price) / self.buy_price
            if current_drop <= -self.p.trigger * (self.add_count + 1):
                if self.add_count < self.p.max_adds:
                    self.add_position()
 
    def initial_buy(self):
        # 计算初始仓位大小 (10%的账户资金)
        size = self.broker.getvalue() * self.p.initial_risk / self.data.close[0]
        self.order = self.buy(size=size)
        self.buy_price = self.data.close[0]
        self.buy_size = size
        self.add_count = 0
        self.stop_price = self.buy_price * (1 - self.p.stop_loss)
        self.log(f'初始买入 @ {self.data.close[0]:.2f}, 数量: {size:.2f}')
 
    def add_position(self):
        # 计算加仓数量 (每次增加4%的账户资金)
        add_size = self.broker.getvalue() * self.p.add_risk / self.data.close[0]
        total_size = self.buy_size * (1 + self.add_count * (self.p.add_risk / self.p.initial_risk)) + add_size
        
        # 检查最大加仓次数
        if self.add_count >= self.p.max_adds:
            self.log(f'已达到最大加仓次数 {self.p.max_adds}')
            return
            
        # 执行加仓
        self.order = self.buy(size=add_size)
        self.buy_size = total_size
        self.add_count += 1
        self.buy_price = (self.buy_price * self.buy_size - add_size * self.data.close[0]) / (self.buy_size - add_size)
        self.stop_price = self.buy_price * (1 - self.p.stop_loss)
        self.log(f'第 {self.add_count} 次加仓 @ {self.data.close[0]:.2f}, 加仓数量: {add_size:.2f}, 平均价格: {self.buy_price:.2f}')
    
    # 止损平仓
    def close_all_positions(self,text='止损'):
        if self.position:
            self.order = self.close()
            self.log(f'{text}平仓 @ {self.data.close[0]:.2f}')
            self.add_count = 0
            self.buy_price = 0.0
            self.buy_size = 0.0


    # def log(self, txt, dt=None):
    #     dt = dt or self.datas[0].datetime.date(0)
    #     print(f'{dt.isoformat()}, {txt}')
 
    def stop(self):
        print(f'初始资金: {self.broker.startingcash:.2f}')
        print(f'最终资金: {self.broker.getvalue():.2f}')
        print(f'总盈亏: {self.broker.getvalue() - self.broker.startingcash:.2f}')